import os
import re
from wordcloud import WordCloud
import matplotlib.pyplot as plt
from collections import Counter
import csv

output_dir = 'output'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)
    print(f"Created directory: {output_dir}")
# 第一部分：处理文件并保存到logs目录
def process_files_and_save():
    # 创建logs目录（如果不存在）
    logs_dir = "logs"
    if not os.path.exists(logs_dir):
        os.makedirs(logs_dir)
        print(f"Created directory: {logs_dir}")

    # 遍历当前目录下的文件
    files_found = False
    for filename in os.listdir("."):
        print(f"Checking file: {filename}")  # 调试信息
        # 匹配文件名类似"20250301T230000"的文件
        if re.match(r"[0-9T]", filename):
            files_found = True
            try:
                # 读取文件内容
                with open(filename, "r", encoding="utf-8") as file:
                    content = file.read()
                    print(f"Content of {filename}: {content[:100]}...")  # 调试信息

                # 将制表符替换为逗号
                content = content.replace("\t", ",")
                print(f"Processed content: {content[:100]}...")  # 调试信息

                # 保存到logs目录下，文件名不变，扩展名为.csv
                new_filename = os.path.join(logs_dir, f"{filename}.csv")
                with open(new_filename, "w", encoding="utf-8") as new_file:
                    new_file.write(content)

                print(f"Processed and saved: {new_filename}")
            except Exception as e:
                print(f"Error processing file {filename}: {e}")

    if not files_found:
        print("No files found matching the pattern '^\d{14}$'. Please check the file names.")

# 第二部分：从logs目录遍历文件，按列统计并生成词云
def generate_word_cloud_from_logs():
    logs_dir = "logs"
    if not os.path.exists(logs_dir):
        print("No 'logs' directory found. Please run the first part first.")
        return

    # 遍历logs目录下的所有.csv文件
    files_found = False
    processed_files = set()  # 用于记录已生成词云的文件
    for filename in os.listdir(logs_dir):
        if filename.endswith(".csv"):
            files_found = True
            if filename in processed_files:
                print(f"Skipping already processed file: {filename}")
                continue
            file_path = os.path.join(logs_dir, filename)
            print(f"Processing CSV file: {file_path}")  # 调试信息
            # 分析生成的 CSV 文件，按第五列（设备名称）统计结果
            device_names = []
            ignored_devices = {"小米通用协议", "DoT", "HTTPS", "SSL", "DoH"}  # 定义需要忽略的设备名称集合
            try:
                with open(file_path, "r", encoding="utf-8") as infile:
                    reader = csv.reader(infile)
                    next(reader)  # 跳过标题行（如果有）
                    for row in reader:
                        if len(row) >= 5:  # 确保行有足够多的列
                            device_name = row[4]  # 假设第五列是设备名称
                            if device_name not in ignored_devices:  # 如果设备名称不在忽略列表中
                                device_names.append(device_name)  # 添加到设备名称列表

                # 将设备名称列表转换为字符串
                all_text = " ".join(device_names)

                # 设置中文字体（确保支持中文显示）
                font_path = "simhei.ttf"  # 黑体字体路径，可根据实际情况调整
                if not os.path.exists(font_path):
                    print(f"字体文件 {font_path} 未找到，请确保字体文件存在！")
                    exit()

                # 生成词云
                if all_text.strip():  # 确保内容不为空
                    wordcloud = WordCloud(font_path=font_path,width=1920, height=1080, background_color="white").generate(all_text)
                    

                    # 保存词云图片
                    output_image_path = os.path.join(output_dir, f"{filename[:-4]}_wordcloud.png")
                    wordcloud.to_file(output_image_path)

                    print(f"Generated word cloud for {filename} and saved as {output_image_path}")
                    processed_files.add(filename)  # 记录已生成词云的文件
                else:
                    print(f"No text content found in {filename}. Skipping word cloud generation.")
            except Exception as e:
                print(f"Error processing file {filename}: {e}")


    if not files_found:
        print("No .csv files found in the 'logs' directory. Please check the files.")

# 执行第一部分：处理文件并保存到logs目录
process_files_and_save()

# 执行第二部分：从logs目录生成词云
generate_word_cloud_from_logs()
